import numpy as np
import struct
import argparse
import os

def read_int8_matrix(filename, rows, cols):
    if not os.path.exists(filename):
        print(f"错误: 文件未找到 -> {filename}")
        return None
    with open(filename, 'r') as f:
        hex_lines = [line.strip() for line in f if line.strip()]
        int8_values = []
        for hex_str in hex_lines:
            byte = int(hex_str, 16)
            if byte >= 0x80:
                byte -= 0x100
            int8_values.append(byte)
        return np.array(int8_values, dtype=np.int8).reshape(rows, cols)

def read_int8_matrix_col_major(filename, rows, cols):
    if not os.path.exists(filename):
        print(f"错误: 文件未找到 -> {filename}")
        return None
    with open(filename, 'r') as f:
        hex_lines = [line.strip() for line in f if line.strip()]
        int8_values = []
        for hex_str in hex_lines:
            byte = int(hex_str, 16)
            if byte >= 0x80:
                byte -= 0x100
            int8_values.append(byte)
        return np.array(int8_values, dtype=np.int8).reshape(rows, cols, order='F')

def read_int4_matrix(filename, rows, cols):
    if not os.path.exists(filename):
        print(f"错误: 文件未找到 -> {filename}")
        return None
    with open(filename, 'r') as f:
        hex_lines = [line.strip() for line in f if line.strip()]
        int4_values = []
        for hex_str in hex_lines:
            val = int(hex_str, 16)
            # int4有符号补码，0x0~0x7为0~7，0x8~0xF为-8~-1
            if val >= 0x8:
                val -= 0x10
            int4_values.append(val)
        return np.array(int4_values, dtype=np.int8).reshape(rows, cols)

def read_int4_matrix_col_major(filename, rows, cols):
    if not os.path.exists(filename):
        print(f"错误: 文件未找到 -> {filename}")
        return None
    with open(filename, 'r') as f:
        hex_lines = [line.strip() for line in f if line.strip()]
        int4_values = []
        for hex_str in hex_lines:
            val = int(hex_str, 16)
            if val >= 0x8:
                val -= 0x10
            int4_values.append(val)
        return np.array(int4_values, dtype=np.int8).reshape(rows, cols, order='F')

def read_int32_matrix_from_big_endian_hex(filename, rows, cols):
    if not os.path.exists(filename):
        print(f"错误: 文件未找到 -> {filename}")
        return None
    with open(filename, 'r') as f:
        hex_lines = [line.strip() for line in f if line.strip()]
        int32_values = []
        for hex_str in hex_lines:
            int32_val = int(hex_str, 16)
            if int32_val >= 0x80000000:
                int32_val = int32_val - 0x100000000
            int32_values.append(int32_val)
        return np.array(int32_values, dtype=np.int32).reshape(rows, cols)

def saturating_add_int32(a, b):
    """
    对两个int32的numpy数组执行饱和加法，手动判断溢出
    当结果大于最大int32时，表示为int32最大值；小于最小值则表示为最小值
    """
    min_val = np.iinfo(np.int32).min
    max_val = np.iinfo(np.int32).max
    
    # 手动判断溢出
    result = np.zeros_like(a, dtype=np.int32)
    
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            # 获取当前元素
            a_val = a[i, j]
            b_val = b[i, j]
            
            # 手动判断溢出
            if a_val > 0 and b_val > 0:
                # 两个正数相加，检查是否溢出
                if a_val > max_val - b_val:
                    result[i, j] = max_val  # 溢出，设为最大值
                else:
                    result[i, j] = a_val + b_val
            elif a_val < 0 and b_val < 0:
                # 两个负数相加，检查是否溢出
                if a_val < min_val - b_val:
                    result[i, j] = min_val  # 溢出，设为最小值
                else:
                    result[i, j] = a_val + b_val
            else:
                # 一正一负，不会溢出
                result[i, j] = a_val + b_val
    
    return result

def main():
    parser = argparse.ArgumentParser(description='可配置的矩阵乘加模拟脚本，支持int8/int4')
    parser.add_argument('--file_a', required=True, help='A矩阵文件路径')
    parser.add_argument('--file_b', required=True, help='B矩阵文件路径')
    parser.add_argument('--file_c', required=True, help='C矩阵文件路径')
    parser.add_argument('--output_file', required=True, help='输出文件路径')
    parser.add_argument('--m_size', type=int, required=True, help='A的行数')
    parser.add_argument('--n_size', type=int, required=True, help='B的列数')
    parser.add_argument('--k_size', type=int, required=True, help='A的列数/B的行数')
    parser.add_argument('--dtype', choices=['int8', 'int4'], required=True, help='输入数据类型')
    parser.add_argument('--b_col_major', action='store_true', help='B矩阵是否列优先')
    args = parser.parse_args()

    print("开始执行硬件行为模拟 (可配置参数)...")
    print(f"A: {args.file_a}\nB: {args.file_b}\nC: {args.file_c}\n输出: {args.output_file}")
    print(f"尺寸: M={args.m_size}, N={args.n_size}, K={args.k_size}, 类型={args.dtype}, B列优先={args.b_col_major}")

    # 读取A
    if args.dtype == 'int8':
        A = read_int8_matrix(args.file_a, args.m_size, args.k_size)
    else:
        A = read_int4_matrix(args.file_a, args.m_size, args.k_size)
    # 读取B
    if args.dtype == 'int8':
        if args.b_col_major:
            B = read_int8_matrix_col_major(args.file_b, args.k_size, args.n_size)
        else:
            B = read_int8_matrix(args.file_b, args.k_size, args.n_size)
    else:
        if args.b_col_major:
            B = read_int4_matrix_col_major(args.file_b, args.k_size, args.n_size)
        else:
            B = read_int4_matrix(args.file_b, args.k_size, args.n_size)
    # 读取C
    C = read_int32_matrix_from_big_endian_hex(args.file_c, args.m_size, args.n_size)

    if A is None or B is None or C is None:
        print("错误: 一个或多个输入文件无法读取，程序终止。")
        return

    print("正在计算矩阵乘法 (A * B)...")
    print(f"A矩阵形状: {A.shape}")
    print(f"B矩阵形状: {B.shape}")
    AB = np.matmul(A.astype(np.int32), B.astype(np.int32), dtype=np.int32)
    print(f"AB矩阵形状: {AB.shape}")

    print("正在执行饱和加法 (AB + C)...")
    print(f"C矩阵形状: {C.shape}")
    result = saturating_add_int32(AB, C)
    print(f"结果矩阵形状: {result.shape}")

    print(f"正在将结果写入文件: {args.output_file}")
    with open(args.output_file, 'w') as f:
        for row in result:
            for val in row:
                packed = struct.pack('>i', val)
                hex_str = packed.hex()
                f.write(hex_str + '\n')

    print("\n模拟完成！")
    print(f"生成的黄金参考模型 (Golden Model) 已保存至: {args.output_file}")
    print(f"结果矩阵尺寸: {args.m_size}x{args.n_size} = {args.m_size*args.n_size} 个元素")

if __name__ == '__main__':
    main() 